# -*- coding: utf-8 -*-
"""
Created on Mon Sep 12 13:48:33 2022

@author: Shang'xiang

Stay Hungry Stay Foolish
"""

'''
此程序用于走时CT的GN反演
'''

# 初始化，读取数据
import ttcrpy.rgrid as rg
import numpy as np
from scipy.sparse.linalg import cg

import matplotlib.pyplot as plt

data = np.loadtxt('G03_FSM.txt')

srcs = data[:,0:2]
rcvs = data[:,2:4]
tobs = data[:,4]

# 给定初始模型
# 创建网格
x = np.arange(0,31.0)
y = np.arange(0,21.0)

# 创建速度模型
v = 2000*np.ones((x.size-1,y.size-1))


# 离散网格
grid = rg.Grid2d(x, y, cell_slowness=True, method = 'SPM')

# 速度转换为慢度
slowness = 1./v

# =============================================================================
# # tcal,LL = grid.raytrace(srcs, rcvs, slowness, compute_L = True)
# 
# # A = LL.todense()
# 
# # I = np.eye(600)
# 
# # zuo = A.T*A + 0.618*I
# # deltat = tobs - tcal
# # you = A.T.dot(deltat)
# # you = you.T
# 
# # deltam,info = cg(zuo,you)
# =============================================================================

I = np.eye(600)


def s2dLaplaceEqual(nex,ney):
# =============================================================================
#     s2dLaplace: equal grid laplace algorithm
# =============================================================================
    gm = np.zeros([nex*ney,nex*ney])
    jj = -1
    for ix in range(nex):
        for iy in range(ney):
            jj = jj + 1
            c0 = 0.0
            # left
            if ix > 0:
                if iy >= 0:
                    t = 1
                    c0 = c0 + t
                    ii = jj - ney
                    gm[jj,ii] = -t
                    gm[ii,jj] = -t
            # top
            if iy > 0:
                t = 1
                c0 = c0 + t
                ii = jj - 1
                gm[jj,ii] = -t
                gm[ii,jj] = -t
            # bottom
            if iy < ney - 1:
                t = 1
                c0 = c0 + t
                ii = jj + 1
                gm[jj,ii] = -t
                gm[ii,jj] = -t
            # right
            if ix < nex - 1:
                if iy >= 0:
                    t = 1
                    c0 = c0 + t
                    ii = jj + ney
                    gm[jj,ii] = -t
                    gm[ii,jj] = -t
            gm[jj,jj] = c0
    return gm

C = s2dLaplaceEqual(30,20)
            

# =============================================================================
# # CT正演
# =============================================================================
def ctforward(slowness):
    tcal,LL = grid.raytrace(srcs, rcvs, slowness, compute_L = True)
    A = LL.todense()
    zuo = A.T*A + 0.618*I
    # zuo = A.T*A + 0.618*C
    deltat = tobs - tcal
    you = A.T.dot(deltat)
    you = you.T
    deltam,info = cg(zuo,you)
    return deltam,deltat

# 第一次计算
# 计算模型修改量
delta = ctforward(slowness)
deltam = delta[0]
deltam = deltam.reshape(30,20)
# 
# # 修改慢度
slowness = slowness + deltam

# 开始循环
rms_all = np.empty(0)
for i in range(10):
    delta = ctforward(slowness)
    deltam = delta[0]
    deltat = delta[1]
    rms = np.sum(deltat/tobs)
    rms_all = np.append(rms_all, rms)
    deltam = deltam.reshape(30,20)
    slowness = slowness + deltam
    v = 1./slowness
    # 画模型图
    fig, ax = plt.subplots()
    cs = plt.pcolor(v,cmap='jet',edgecolor='w')
    plt.colorbar()
    plt.show()
    
# 保存误差
np.savetxt('GN_rms.txt', rms_all)